package com.lihainuo.springai.redis.strater;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 自定义 Redis 存储 Memory
 * Message -> SerializableMessage
 * ||                          ||
 * SerializableMessage -> Message
 */
public class RedisChatMemory implements ChatMemory {

    private final RedisTemplate<String, SerializableMessage> redisTemplate;

    /**
     * 构造器注入
     * @param redisTemplate Redis 模板
     */
    public RedisChatMemory(RedisTemplate<String, SerializableMessage> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public void add(String conversationId, Message message) {
        // 转换数据
        SerializableMessage serializableMessage = new SerializableMessage(message);
        // 向 Redis 列表中添加数据 左侧添加
        redisTemplate.opsForList().leftPush(conversationId, serializableMessage);
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // 转换数据
        List<SerializableMessage> serializableMessages = messages.stream()
                .map(SerializableMessage::new)
                .collect(Collectors.toList());
        // 向 Redis 列表中添加数据 左侧添加
        if (serializableMessages != null && !serializableMessages.isEmpty()) {
            redisTemplate.opsForList().leftPushAll(conversationId, serializableMessages);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 需要将 SerializableMessage -> Message
        List<SerializableMessage> serializableMessages = redisTemplate.opsForList()
                .range(conversationId, 0, lastN - 1);
        // 处理空结果
        if (serializableMessages == null || serializableMessages.isEmpty()) {
            return List.of();
        }
        // 处理非空结果
        return serializableMessages.stream()
                .map(SerializableMessage::toMessage)
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        // 清空 Redis 列表
        redisTemplate.delete(conversationId);
    }
}
